//
// Created by 邹迪凯 on 2021/7/17.
//

#ifndef LWPT_PUBLICRANDOMTASKASSIGNMENT_H
#define LWPT_PUBLICRANDOMTASKASSIGNMENT_H

#include "assignment.h"

class PRTAServer;

class PRTAWorker;

class PRTAApplicationChunk : public Chunk {
public:
    string id;
    string lat;
    string lon;

    PRTAApplicationChunk(string id, string lat, string lon) {
        this->id = id;
        this->lat = lat;
        this->lon = lon;
    }

    long size() const override {
        return id.size() + lat.size() + lon.size() + overhead();
    }

    long overhead() const {
        return 1;
    }
};

class GenericRegisterChunk : public Chunk {
public:
    string id;
    string lat;
    string lon;

    GenericRegisterChunk(const Task &task) {
        id = task.id;
        lat = task.location.lat;
        lon = task.location.lon;
    }

    long size() const override {
        return id.size() + lat.size() + lon.size() + overhead();
    }

    long overhead() const {
        return 1;
    }
};

class PRTAWorker : public Worker {
public:
    PRTAWorker(string id, Location location) : Worker(id, location) {}

    Chunk *apply(Server &server) override;
};

class PRTAServer : public Server {
public:
    vector<Task> tasks;
    vector<PRTAWorker> workers;
    char centric;
    double candidateCons = 100000;

    unordered_map<string, unordered_map<string, double> *> taskWorkerDistance;
    set<string> candidateTasks, candidateWorkers;
    unordered_map<string, double> wx, wy;
    unordered_map<string, string> cx, cy;
    unordered_map<string, bool> visx, visy;
    double delta;
    double gb = 2000;

    PRTAServer(char centric) : centric(centric) {}

    void prepareForAssignment() override {
        for (const auto &reg : *registrations) {
            auto r = static_cast<GenericRegisterChunk *>(reg);
            tasks.emplace_back(r->id, Location(r->lat, r->lon));
        }
        for (const auto &ap : *applications) {
            auto a = static_cast<PRTAApplicationChunk *>(ap);
            workers.emplace_back(a->id, Location(a->lat, a->lon));
        }
    }

    void assignTask() override {
        if (centric == 'p') {
            assignPCentric();
        } else if (centric == 't') {
            assignTCentric();
        } else {
            assignKM();
        }
    }

    void assignKM() {
        for (const auto &item : workers) candidateWorkers.insert(item.id);
        for (const auto &item : tasks) candidateTasks.insert(item.id);
        for (auto &task:tasks) {
            if (taskWorkerDistance.find(task.id) == taskWorkerDistance.end()) {
                taskWorkerDistance[task.id] = new unordered_map<string, double>;
            }
            for (auto &worker:workers) {
                auto tl = task.location;
                auto wl = worker.location;
                (*taskWorkerDistance[task.id])[worker.id] = earthDistance(stod(tl.lat), stod(tl.lon), stod(wl.lat),
                                                                          stod(wl.lon));
            }
        }
        //把矩阵补足成方阵
        if (candidateWorkers.size() < candidateTasks.size()) {
            int dummyCount = candidateTasks.size() - candidateWorkers.size();
            for (int i = 0; i < dummyCount; i++) {
                string dummyWorkerId = generateId(ID_LENGTH, candidateWorkers);
                candidateWorkers.insert(dummyWorkerId);
                for (const auto &task : taskWorkerDistance) {
                    (*task.second)[dummyWorkerId] = 2 * gb;
                }
            }
        } else if (candidateWorkers.size() > candidateTasks.size()) {
            int dummyCount = candidateWorkers.size() - candidateTasks.size();
            for (int i = 0; i < dummyCount; i++) {
                string dummyTaskId = generateId(ID_LENGTH, candidateTasks);
                candidateTasks.insert(dummyTaskId);
                taskWorkerDistance[dummyTaskId] = new unordered_map<string, double>;
                for (const auto &workerId : candidateWorkers) {
                    (*taskWorkerDistance[dummyTaskId])[workerId] = 2 * gb;
                }
            }
        }
        for (auto &taskWorkers:taskWorkerDistance) {
            for (auto &cw:candidateWorkers) {
                if (taskWorkers.second->find(cw) == taskWorkers.second->end()) {
                    (*taskWorkers.second)[cw] = 2 * gb;
                }
            }
        }
        for (auto &taskWorkers:taskWorkerDistance) {
            for (auto &workerDistance:*taskWorkers.second) {
                workerDistance.second = -workerDistance.second;
            }
        }

        km();

        for (auto &taskWorkers:taskWorkerDistance) {
            auto taskId = taskWorkers.first;
            if (!cx[taskId].empty() &&
                !equals(taskWorkerDistance[taskId]->at(cx[taskId]), -2 * gb)) {
                assignments.push_back(new pair<string, string>(taskId, cx[taskId]));
                assignedTaskCount++;
            }
        }
    }

    bool dfs(string tt) {
        visx[tt] = true;
        for (const auto &ww:candidateWorkers) {
            if (!visy[ww]) {
                double t = wx[tt] + wy[ww] - taskWorkerDistance[tt]->at(ww);
                if (equals(t, 0)) {
                    visy[ww] = true;
                    if (cy[ww].empty() || dfs(cy[ww])) {
                        cx[tt] = ww;
                        cy[ww] = tt;
                        return true;
                    }
                } else if (t > 0) {
                    delta = min(delta, t);
                }
            }
        }
        return false;
    }

    void km() {
        const string EMPTY_STR;
        memset(cx, candidateTasks, EMPTY_STR);
        memset(cy, candidateWorkers, EMPTY_STR);
        memset(wx, candidateTasks, 0.0);
        memset(wy, candidateWorkers, 0.0);
        for (const auto &t:candidateTasks) {
            for (const auto &w:candidateWorkers) {
                wx[t] = max(wx[t], taskWorkerDistance[t]->at(w));
            }
        }
        for (const auto &t:candidateTasks) {
            while (true) {
                delta = INF;
                memset(visx, candidateTasks, false);
                memset(visy, candidateWorkers, false);
                if (dfs(t))
                    break;
                for (const auto &tt:candidateTasks)
                    if (visx[tt])
                        wx[tt] -= delta;
                for (const auto &ww:candidateWorkers)
                    if (visy[ww])
                        wy[ww] += delta;
            }
        }
    }

    void assignPCentric() {
        while (!workers.empty()) {
            auto workerIndex = rand() % workers.size();
            auto worker = workers[workerIndex];
            set<int> candidates;
            for (int i = 0; i < tasks.size(); i++) {
                if (worker.location.distance(tasks[i].location) < candidateCons) {
                    candidates.insert(i);
                }
            }
            if (!candidates.empty()) {
                assignedTaskCount++;
                double minDistance = 100000;
                int minDistanceTask = 0;
                for (auto ti:candidates) {
                    double distance = tasks[ti].location.distance(worker.location);
                    if (distance < minDistance) {
                        minDistanceTask = ti;
                        minDistance = distance;
                    }
                }
                assignments.push_back(new pair<string, string>(tasks[minDistanceTask].id, worker.id));
                swap(tasks, minDistanceTask, tasks.size() - 1);
                tasks.pop_back();
            }
            swap(workers, workerIndex, workers.size() - 1);
            workers.pop_back();
        }
    }

    void assignTCentric() {
        while (!tasks.empty()) {
            auto taskIndex = rand() % tasks.size();
            auto task = tasks[taskIndex];
            set<int> candidates;
            for (int i = 0; i < workers.size(); i++) {
                if (task.location.distance(workers[i].location) < candidateCons) {
                    candidates.insert(i);
                }
            }
            if (!candidates.empty()) {
                assignedTaskCount++;
                double minDistance = 100000;
                int minDistanceWorker = 0;
                for (auto wi:candidates) {
                    double distance = workers[wi].location.distance(task.location);
                    if (distance < minDistance) {
                        minDistanceWorker = wi;
                        minDistance = distance;
                    }
                }
                assignments.push_back(new pair<string, string>(task.id, workers[minDistanceWorker].id));
                swap(workers, minDistanceWorker, workers.size() - 1);
                workers.pop_back();
            }
            swap(tasks, taskIndex, tasks.size() - 1);
            tasks.pop_back();
        }
    }
};

Chunk *PRTAWorker::apply(Server &server) {
    return new PRTAApplicationChunk(id, location.lat, location.lon);
}

class PRTATest : public Test {
public:
    char centric;

    PRTATest(DataSource *ds, int seed, char centric = 'k') : Test(ds), centric(centric) {
        srand(seed);
    }

    Chunk *createTaskRegistration(Task &task) override {
        return new GenericRegisterChunk(task);
    }

    Server *serverFactory() override {
        return new PRTAServer(centric);
    }

    Worker *workerFactory(string id, Location location) override {
        return new PRTAWorker(id, location);
    }

    string name() override {
        return string("PRTA(").append(&centric).append(")");
    }
};

#endif //LWPT_PUBLICRANDOMTASKASSIGNMENT_H
